from entity.train_data import TrainData
from log import log
from util.sqlmapper import SQLManager


class TrainDataDao:
    def getAllIId(self):
        with SQLManager() as db:
            return db.get_all("select id from t_train")

    def deleteById(self, train_id: int):
        with SQLManager() as db:
            db.modify("delete from t_train where id = %s", train_id)

    def getAllEACCImgName(self):
        with SQLManager() as db:
            return db.get_all("select epoch_acc_img_name from t_train where epoch_acc_img_name is not null")
    def getAllModelImgName(self):
        with SQLManager() as db:
            return db.get_all("select model_img_name from t_train")

    def getEACCImgNameByTrainId(self, train_id):
        with SQLManager() as db:
            msg = db.get_one("select epoch_acc_img_name from t_train where id = %s", train_id)
            if msg is None:
                return msg
            return msg[0]
    def getModelImgNameByTrainId(self, train_id):
        with SQLManager() as db:
            msg = db.get_one("select model_img_name from t_train where id = %s", train_id)
            if msg is None:
                return msg
            return msg[0]

    def setExcludeImgName(self, t: TrainData, save_path):
        with SQLManager() as db:
            log.info("insert t in db")
            return db.insert_return_key("insert into t_train values (null,%s,%s,%s,%s,%s, null,%s)",
                                        t.count_epoch,
                                        t.learning_rate,
                                        t.decay,
                                        t.image_height,
                                        t.image_width,
                                        save_path
                                        )

    def findById(self, train_id: int) -> TrainData:
        with SQLManager() as db:
            log.info("insert t in db")
            res = db.get_one("select * from t_train where id = %s", train_id)
            t = TrainData()
            t.count_epoch = res[1]
            t.learning_rate = res[2]
            t.decay = res[3]
            t.image_height = res[4]
            t.image_width = res[5]
            t.epoch_acc_img_name = res[6]
            t.id = train_id
            return t

    def setImgNameCountEpochById(self, train_id: int, img_name: str, current_epoch: int):
        with SQLManager() as db:
            log.info("set imgName in db")
            db.modify("update t_train set epoch_acc_img_name = %s,count_epoch= %s where id = %s",
                      img_name,
                      current_epoch,
                      train_id
                      )

    def setModelImgNameById(self, train_id: int, name: str):
        with SQLManager() as db:
            log.debug("setModelImgNameById")
            db.modify("update t_train set model_img_name = %s where id = %s", name, train_id)


if __name__ == "__main__":
    # print(TrainDataDao().findById(14))
    print(type(TrainDataDao().getEACCImgNameByTrainId(56)))
